from src.测量比例尺 import detect_scale_bar
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt


def _limit_long_side(img, max_long=1000):
    """将长边限制在 max_long 以内，返回缩放后的图像及比例"""
    h, w = img.shape[:2]
    scale = min(max_long / max(h, w), 1.0)  # 不放大
    if scale < 1.0:
        new_size = (int(w * scale), int(h * scale))
        return cv2.resize(img, new_size), scale
    return img.copy(), 1.0


def _put_text_pil(img_bgr, text, org=(10, 30),
                  font_path="simsun.ttc", font_size=24, color=(255, 255, 255)):
    """使用 PIL 在 OpenCV 图像上写中文"""
    img_pil = Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    font = ImageFont.truetype(font_path, font_size)
    draw.text(org, text, font=font, fill=color[::-1])  # RGB -> BGR
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)


def _select_roi_scaled(window_name, img_scaled, scale):
    """在窗口中让用户框选 ROI，返回相对于原图的 (x, y, w, h)"""
    clone = img_scaled.copy()
    roi_scaled = None
    selecting = True

    def on_mouse(event, x, y, flags, param):
        nonlocal roi_scaled, selecting
        if event == cv2.EVENT_LBUTTONDOWN:
            roi_scaled = [x, y, 0, 0]
        elif event == cv2.EVENT_MOUSEMOVE and roi_scaled is not None:
            roi_scaled[2] = x - roi_scaled[0]
            roi_scaled[3] = y - roi_scaled[1]
            tmp = clone.copy()
            cv2.rectangle(tmp,
                          (roi_scaled[0], roi_scaled[1]),
                          (x, y),
                          (0, 255, 0), 2)
            cv2.imshow(window_name, tmp)
        elif event == cv2.EVENT_LBUTTONUP:
            selecting = False

    cv2.imshow(window_name, clone)
    cv2.setMouseCallback(window_name, on_mouse)

    while selecting or roi_scaled is None:
        key = cv2.waitKey(30) & 0xFF
        if key == 27:  # Esc 取消
            cv2.destroyWindow(window_name)
            return None

    cv2.destroyWindow(window_name)

    # 将缩放坐标恢复为原图坐标
    x, y, w, h = [int(v / scale) for v in roi_scaled]
    return x, y, w, h


def manual_box_selection(img, max_long=1000, win_max=(800, 600), font_path="simsun.ttc"):
    """
    手动框选条纹区域
    :param img:         原始图像（OpenCV BGR ndarray）
    :param max_long:    长边限制，超过则等比缩小用于显示
    :param win_max:     窗口最大宽高 (w, h)
    :param font_path:   中文字体路径（Windows 可用 'simsun.ttc'，Linux/Mac 可换其他）
    :return:            (roi_img, roi_on_original) 或 (None, None)
    """
    # 1. 统一缩放
    img_scaled, scale = _limit_long_side(img, max_long)

    # 2. 绘制提示文字
    tip = "请框选条纹区域，按 Enter 确认，Esc 取消"
    img_display = _put_text_pil(img_scaled, tip, org=(10, 30),
                                font_path=font_path, font_size=22, color=(255, 255, 255))

    # 3. 创建窗口并限制大小
    h, w = img_scaled.shape[:2]
    win_scale = min(win_max[0] / w, win_max[1] / h, 1.0)
    win_size = (int(w * win_scale), int(h * win_scale))

    cv2.namedWindow("框选条纹区域", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("框选条纹区域", *win_size)
    cv2.imshow("框选条纹区域", img_display)

    # 4. 交互式选取 ROI
    roi = _select_roi_scaled("框选条纹区域", img_scaled, scale)
    if roi is None:
        return None, None

    x, y, w, h = roi
    roi_img = img[y:y + h, x:x + w]
    return roi_img, (x, y, w, h)


def auto_dense_stripe_roi(img, win_size=16, stride=4, thres_percentile=99.5, margin=1):
    """
    自动检测图像中最密集条纹区域（梯度检测），返回 ROI 框 (x, y, w, h)。
    参数:
        img : numpy.ndarray
            8-bit 灰度图像
        win_size : int, 可选
            滑动窗口大小，用于局部特征分析
        stride : int, 可选
            滑动窗口步长
        thres_percentile : int, 可选
            自适应阈值百分位数，用于筛选高得分区域
        margin : int, 可选
            ROI边界扩展像素数
    返回:
        tuple 或 None
            ROI框坐标 (x, y, width, height)，若未检测到则返回None
    """
    # 获取图像尺寸
    h, w = img.shape[:2]

    # 1) 计算梯度能量图 - 条纹区域通常具有较高的梯度变化
    gx = cv2.Sobel(img, cv2.CV_64F, 1, 0, ksize=3)  # x方向梯度
    gy = cv2.Sobel(img, cv2.CV_64F, 0, 1, ksize=3)  # y方向梯度
    grad_mag = np.sqrt(gx ** 2 + gy ** 2)  # 梯度幅值

    # 2) 滑动窗口统计每个块的能量特征
    # 生成滑动窗口的左上角坐标网格
    idx_y, idx_x = np.mgrid[0:h - win_size:stride, 0:w - win_size:stride]
    idx_y = idx_y.ravel()  # 展平为一维数组
    idx_x = idx_x.ravel()
    score = np.empty(len(idx_x))  # 存储每个窗口的评分

    # 遍历所有窗口位置
    for k, (y0, x0) in enumerate(zip(idx_y, idx_x)):
        patch = grad_mag[y0:y0 + win_size, x0:x0 + win_size]
        # 使用均值×方差作为评分指标，同时捕获强度和变化性
        score[k] = np.mean(patch) * np.std(patch)

    # 3) 自适应阈值筛选高得分区域
    th = np.percentile(score, thres_percentile)  # 计算阈值
    mask = score >= th  # 生成掩码
    if not np.any(mask):  # 检查是否有区域满足阈值
        print("未检测到明显的条纹密集区域")
        return None

    # 4) 计算候选区域的最小外接矩形
    # 获取满足条件窗口的中心点坐标
    centers = np.column_stack((idx_x[mask] + win_size // 2,
                               idx_y[mask] + win_size // 2))
    # 计算边界坐标
    x_min, y_min = centers.min(axis=0)
    x_max, y_max = centers.max(axis=0)

    # 5) 添加边界并限制在图像范围内
    x = max(0, int(x_min - margin))  # 左边界
    y = max(0, int(y_min - margin))  # 上边界
    roi_w = min(w - x, int(x_max - x_min + 2 * margin))  # 宽度
    roi_h = min(h - y, int(y_max - y_min + 2 * margin))  # 高度

    roi = (x, y, roi_w, roi_h)  # 组装ROI元组

    # 6) 可视化检测结果
    vis = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)  # 转换为BGR用于绘制彩色矩形
    cv2.rectangle(vis, (x, y), (x + roi_w, y + roi_h), (0, 0, 255), 2)  # 绘制红色矩形
    plt.figure(figsize=(5, 5))
    plt.title('Auto-detected dense stripe ROI')
    plt.imshow(vis[..., ::-1])  # 转换为RGB显示
    plt.axis('off')
    plt.show()

    x, y, w, h = roi
    roi_img = img[y:y + h, x:x + w]

    return roi_img, roi

def fft_stripe_roi(img,
                   high_ratio=0.15,  # 高频环内径（相对中心距离）
                   thres_pct=95,  # 响应图阈值（百分位）
                   margin=0):
    """
    基于傅里叶高频能量自动提取条纹 ROI。
    返回: (x, y, w, h) or None
    """
    if img.ndim == 3:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    else:
        gray = img.copy()

    h, w = gray.shape
    # 1) 补零到方形
    side = max(h, w)
    pad_y = (side - h) // 2
    pad_x = (side - w) // 2
    padded = cv2.copyMakeBorder(gray, pad_y, side - h - pad_y,
                                pad_x, side - w - pad_x,
                                cv2.BORDER_REFLECT_101)

    # 2) FFT
    f = np.fft.fft2(padded.astype(np.float32))
    fshift = np.fft.fftshift(f)
    magnitude = np.abs(fshift)

    # 3) 高频环掩膜
    cy = cx = side // 2
    Y, X = np.ogrid[:side, :side]
    dist = np.sqrt((X - cx) ** 2 + (Y - cy) ** 2)
    max_rad = side / 2 * (1 - high_ratio)
    mask = dist >= max_rad  # True 为高频
    fshift[~mask] = 0  # 抹掉低频
    filtered = np.fft.ifftshift(fshift)
    stripe_map = np.abs(np.fft.ifft2(filtered))  # 条纹响应图

    # 4) 阈值分割
    th = np.percentile(stripe_map, thres_pct)
    _, binary = cv2.threshold(stripe_map.astype(np.uint8), th, 255, cv2.THRESH_BINARY)

    # 5) 最大连通域
    num, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
    if num <= 1:  # 只有背景
        return None
    # stats: [x, y, w, h, area]; 0 号背景跳过
    largest = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
    x, y, bw, bh, _ = stats[largest]

    # 6) 映射回原图坐标
    roi_x = max(0, x - pad_x - margin)
    roi_y = max(0, y - pad_y - margin)
    roi_w = min(w - roi_x, bw + 2 * margin)
    roi_h = min(h - roi_y, bh + 2 * margin)

    return (roi_x, roi_y, roi_w, roi_h)


if __name__ == '__main__':
    # 读图
    img = cv2.imread('../img/new/0012.tif', 0)
    if img is None:
        raise FileNotFoundError('../img/new/0012.tif')
    img, _1, _2 = detect_scale_bar(gray_img=img, img_type=2048)

    # 第一种方法
    img, roi = auto_dense_stripe_roi(img)
    if roi is None:
        print('未检测到条纹区域')
        exit()
    print('自动 ROI:', roi)

    # 第二种方法
    # img = img[roi[1]:roi[1] + roi[3], roi[0]:roi[0] + roi[2]] # 在第一种基础上运行
    # img = img[:1800, :]
    # roi = fft_stripe_roi(img, high_ratio=0.4, thres_pct=90, margin=20)
    # print("ROI:", roi)
    # if roi is not None:
    #     x, y, w, h = roi
    #     vis = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    #     cv2.rectangle(vis, (x, y), (x + w, y + h), (0, 0, 255), 2)
    #     plt.imshow(vis[..., ::-1])
    #     plt.axis('off')
    #     plt.show()
